# -*- coding: utf-8 -*-
# =============================================================================
#     FileName: redis_expire_helper.py
#         Desc: 用于清理Redis服务器上指定前缀的KEY，通过循环SCAN方式降低清理过程对服务器的影响。
#       Author: GGA
#        Email:
#     HomePage:
#      Version: 1.0.1
#   LastChange: 2020-07-31
#      History:
# =============================================================================
import enum
import logging
import multiprocessing
import random
import time

import fire
import redis

from utils.logger_helper import LoggerHelper

app_name = "redis_expire_helper"
app_version = "1.0"
logger = LoggerHelper.get_logger()


class ExpireAction(enum.Enum):
    CleanExpiredKeys = "clean_expired_keys"
    CleanUnexpiredKeys = "clean_unexpired_keys"
    CleanLongExpiredKeys = "clean_long_expired_keys"
    CleanLongIdleKeys = "clean_long_idle_keys"


class RedisExpire(object):
    def __init__(self,
                 expire_action: ExpireAction, key_pattern: str = "demo::*",
                 redis_host: str = "127.0.0.1", redis_port: int = 6379,
                 redis_password: str = "", redis_db: int = 0,
                 scan_size=1000, scan_interval=0.1, pipe_size=100,
                 expire_seconds: int = None, random_seconds: int = None,
                 long_idle_seconds: int = None, log_sample=0.01):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.redis_password = redis_password
        self.redis_db = redis_db
        self.expire_action = expire_action
        self.key_pattern = key_pattern
        self.scan_size = scan_size
        self.scan_interval = scan_interval
        self.pipe_size = pipe_size
        self.long_idle_seconds = long_idle_seconds
        self.expire_seconds = expire_seconds
        self.random_seconds = random_seconds
        self.log_sample = log_sample
        self.redis_pool = redis.ConnectionPool(
            host=self.redis_host,
            port=self.redis_port,
            db=self.redis_db,
            password=self.redis_password,
            max_connections=2
        )

    def _get_redis_connection(self):
        """
        从Redis Pool中获取连接，避免重复建立Redis连接
        :return:
        """
        return redis.Redis(connection_pool=self.redis_pool, decode_responses=True)

    def _check_connection(self):
        redis_conn = self._get_redis_connection()
        try:
            redis_conn.ping()
            return True
        except Exception as _:
            logger.info("无法正常连接Redis，请检查Redis配置信息是否正确。")
            exit(1)

    def _scan_keys(self, start_post=0):
        """
        范围扫描匹配的键值, 并返回匹配的KEYS
        :param start_post:  扫描开始位置点
        :return: stop_post,match_keys
            stop_post: 本次扫描的结束位置
            match_keys：满足条件的KEYS
        """
        redis_conn = self._get_redis_connection()
        stop_post, match_keys = redis_conn.scan(start_post, self.key_pattern, self.scan_size)
        logger.info("在{}:{}实例上扫描起始点:{},扫描结束点:{},扫描数量:{},匹配Key数量:{}".format(
            self.redis_host, self.redis_port,
            start_post, stop_post, self.scan_size,
            len(match_keys)
        ))
        return stop_post, match_keys

    def _get_idle_seconds(self, key_names):
        key_items = dict()
        if len(key_names) == 0:
            return key_items
        redis_conn = self._get_redis_connection()
        redis_pipe = redis_conn.pipeline()
        for key_name in key_names:
            redis_pipe.object("IDLETIME", key=key_name)
        pipe_result = redis_pipe.execute()
        for key_index, key_name in enumerate(key_names):
            key_items[key_name] = pipe_result[key_index]
        return key_items

    def _get_expire_seconds(self, key_names):
        key_items = dict()
        if len(key_names) == 0:
            return key_items
        redis_conn = self._get_redis_connection()
        redis_pipe = redis_conn.pipeline()
        for key_name in key_names:
            redis_pipe.ttl(key_name)
        pipe_result = redis_pipe.execute()
        for key_index, key_name in enumerate(key_names):
            key_items[key_name] = pipe_result[key_index]
        return key_items

    def _expire_keys(self, key_names):
        """
        通过PIPLINE来批量设置指定KEYS的过期时间
        :param key_names:
        :return:
        """
        if len(key_names) == 0:
            return
        redis_conn = self._get_redis_connection()
        redis_pipe = redis_conn.pipeline()
        for key_name in key_names:
            expire_seconds = self.expire_seconds + random.randint(0, self.random_seconds)
            redis_pipe.expire(key_name, expire_seconds)
            if random.random() < self.log_sample:
                logger.info("设置key {} 在{}秒后过期.".format(key_name, expire_seconds))
        redis_pipe.execute()

    def _unlink_keys(self, key_names):
        """
        通过PIPLINE来批量UNLINK指定KEYS的过期时间
        :param key_names:
        :return:
        """
        if len(key_names) == 0:
            return
        redis_conn = self._get_redis_connection()
        redis_pipe = redis_conn.pipeline()
        for key_name in key_names:
            if random.random() < self.log_sample:
                logger.info("直接删除key {} ".format(key_name))
            redis_pipe.unlink(key_name)
        redis_pipe.execute()

    def _get_long_idle_keys(self, key_names):
        """
        通过PIPLINE来批量获取指定KEYS的空闲时间，并返回空闲时间超过阈值的KEYS
        :param key_names:
        :return:
        """
        match_keys = list()
        key_items = self._get_idle_seconds(key_names=key_names)
        for key_name, idle_seconds in key_items.items():
            if idle_seconds is not None and int(idle_seconds) > self.long_idle_seconds:
                if random.random() < self.log_sample:
                    logger.info("键值{}当前的空闲时间为{},超过阈值{}".format(
                        key_name, idle_seconds, self.long_idle_seconds))
                match_keys.append(key_name)
        return match_keys

    def _get_unexpired_keys(self, key_names):
        """
        通过PIPLINE来批量获取指定KEYS的空闲时间，并返回未设置过期时间的KEY.
        :param key_names:
        :return:
        """
        match_keys = list()
        key_items = self._get_expire_seconds(key_names=key_names)
        for key_name, expire_seconds in key_items.items():
            if expire_seconds is not None and expire_seconds == -1:
                if random.random() < self.log_sample:
                    logger.info("键值{}未设置过期时间".format(key_name))
                match_keys.append(key_name)
            pass
        return match_keys

    def _get_long_expired_keys(self, key_names):
        """
        通过PIPLINE来批量获取指定KEYS的过期时间，并返回过期时间超过阈值的KEYS
        :param key_names:
        :return:
        """
        match_keys = list()
        key_items = self._get_expire_seconds(key_names=key_names)
        for key_name, expire_seconds in key_items.items():
            if expire_seconds is not None and expire_seconds > self.expire_seconds:
                if random.random() < self.log_sample:
                    logger.info("键值{}的过期时间为{},超过阈值{}".format(
                        key_name, expire_seconds, self.expire_seconds
                    ))
                match_keys.append(key_name)
            pass
        return match_keys

    def _spilt_keys(self, key_names):
        """
        将KEY按照pipe操作批次拆分多个pipe组
        :param key_names:
        :return:
        """
        pipe_groups = []
        group_keys = []
        for key_name in key_names:
            group_keys.append(key_name)
            if len(group_keys) == self.pipe_size:
                pipe_groups.append(group_keys)
                group_keys = []
            pass
        if len(group_keys) > 0:
            pipe_groups.append(group_keys)
        return pipe_groups

    def _get_need_expired_keys(self, key_names):
        if self.expire_action == ExpireAction.CleanLongIdleKeys:
            return self._get_long_idle_keys(key_names=key_names)
        elif self.expire_action == ExpireAction.CleanLongExpiredKeys:
            return self._get_long_expired_keys(key_names=key_names)
        elif self.expire_action == ExpireAction.CleanUnexpiredKeys:
            return self._get_unexpired_keys(key_names=key_names)
        else:
            return []

    def _check_config(self):
        if self.log_sample > 1:
            logger.info("请合理设key_log_sample值!")
            exit(1)
        if self.expire_action == ExpireAction.CleanExpiredKeys:
            return True
        if self.expire_seconds is None or self.expire_seconds < 0:
            logger.info("请合理设置最小过期时间")
            exit(1)
        if self.random_seconds is None or self.random_seconds < 0:
            logger.info("请合理设置最大随机时间")
            exit(1)
        if self.expire_action == ExpireAction.CleanLongIdleKeys:
            if self.long_idle_seconds is None:
                logger.info("请设置最大空闲时长!")
                exit(1)
            if self.long_idle_seconds < 300:
                logger.info("设置最大空闲时长过短")
                exit(1)
        return True

    def _loop_expire_keys(self):
        logger.info("开始扫描{}:{}实例,匹配模式: {}, 清理模式: {}".format(
            self.redis_host, self.redis_port,
            self.key_pattern, self.expire_action.value
        ))
        start_post = 0
        stop_post = 1
        while stop_post > 0:
            stop_post, match_keys = self._scan_keys(start_post=start_post)
            batch_groups = self._spilt_keys(key_names=match_keys)
            for batch_keys in batch_groups:
                dt1 = time.time()
                need_expire_keys = self._get_need_expired_keys(key_names=batch_keys)
                dt2 = time.time()
                self._expire_keys(key_names=need_expire_keys)
                dt3 = time.time()
                logger.info(
                    "在{}:{}实例上,扫描到 {} 个KEY,匹配到 {} 个KEY,匹配操作耗时 {} ms,清理操作耗时 {} ms".format(
                        self.redis_host, self.redis_port,
                        len(batch_keys), len(need_expire_keys), (dt2 - dt1) * 1000, (dt3 - dt2) * 1000,
                    ))
            start_post = stop_post
            time.sleep(self.scan_interval)
        logger.info("完成扫描{}:{}实例".format(self.redis_host, self.redis_port, ))

    def expire_keys(self):
        self._check_config()
        self._check_connection()
        self._loop_expire_keys()


def get_redis_instances(redis_instances_string: str):
    try:
        redis_instances = list()
        redis_items = redis_instances_string.split(",")
        for redis_item in redis_items:
            if redis_item.find(":") > 0:
                redis_host = redis_item.split(":")[0]
                redis_port = int(redis_item.split(":")[1])
            else:
                redis_host = redis_item
                redis_port = 6379
            redis_instances.append({"redis_host": redis_host, "redis_port": redis_port})
        return redis_instances
    except Exception as _:
        logger.info("解析Redis实例列表出错,请检查redis_instances参数")
        exit(-1)


def main(expire_action: str = ExpireAction.CleanExpiredKeys.value, key_pattern: str = "demo::*",
         redis_instances: str = "127.0.0.1:6379,127.0.0.1:6380",
         redis_password: str = "", redis_db: int = 0,
         scan_size=1000, scan_interval=0.1, pipe_size=100,
         expire_seconds: int = None, random_seconds: int = None,
         long_idle_seconds: int = None, log_sample=0.01,
         with_debug=False):
    """
    Redis key清理脚本。\n
    根据expire_action参数选:
        - clean_expired_keys：清理已经过期Key,通过scan触发Redis自动清理。
        - clean_unexpired_keys：清理未设置过期时间的Key。
        - clean_long_idle_keys: 清理长期空闲的Key,通过参数long_idle_seconds来指定最大空闲时间。
        - clean_long_expired_keys：清理过期较长的Key, 对于过期时间超过参数expire_seconds指定的Key,将其过期时间设置为expire_seconds。
    :param expire_action: 清理方式
    :param key_pattern: key模式
    :param redis_instances: redis实例列表
    :param redis_password: redis实例密码
    :param redis_db: redis数据库
    :param scan_size: 每次扫描操作key数量
    :param scan_interval: 每次扫描间隔,单位为秒
    :param pipe_size: 每次pip操作key数量
    :param expire_seconds: 设置key的过期时间
    :param random_seconds: 设置key的随机时间
    :param long_idle_seconds: 最大空闲时间, 在--expire_action="clean_long_idle_key"下，超过该值的key会被立即删除
    :param log_sample: 日志采样率, 0表示不记录
    :param with_debug: 是否开启debug模式
    :return:
    """
    if with_debug:
        LoggerHelper.init_logger(logger_name=app_name, logger_level=logging.DEBUG)
    else:
        LoggerHelper.init_logger(logger_name=app_name, logger_level=logging.INFO)
    global logger
    logger = LoggerHelper.get_logger(logger_name=app_name)
    redis_instances = get_redis_instances(redis_instances)
    process_list = list()
    for redis_instance in redis_instances:
        redis_expire = RedisExpire(
            expire_action=ExpireAction(expire_action), key_pattern=key_pattern,
            redis_host=redis_instance["redis_host"], redis_port=redis_instance["redis_port"],
            redis_password=redis_password, redis_db=redis_db,
            scan_size=scan_size, scan_interval=scan_interval, pipe_size=pipe_size,
            expire_seconds=expire_seconds, random_seconds=random_seconds,
            long_idle_seconds=long_idle_seconds, log_sample=log_sample
        )
        process_item = multiprocessing.Process(target=redis_expire.expire_keys, args=())
        process_item.daemon = True
        process_list.append(process_item)
        process_item.start()
    running_process_count = len(process_list)
    while running_process_count > 0:
        logger.info("当前有{}个线程正在执行，休眠5秒".format(str(running_process_count)))
        time.sleep(5)
        running_process_count = 0
        for process_item in process_list:
            if process_item.is_alive():
                running_process_count += 1
        pass
    logger.info("执行完成")


if __name__ == '__main__':
    fire.Fire(main)
